#!/usr/bin/env python3
# -*- coding: utf-8 -*-
import sys
import os
import grpc
import time

# 导入 P4Runtime 库
sys.path.append(os.path.join(os.path.dirname(os.path.abspath(__file__)), '../../utils/'))
import p4runtime_lib.bmv2
import p4runtime_lib.helper

def write_ipv4_lpm_rule(p4info_helper, sw, dst_ip_addr, dst_mac_addr, port):
    """
    下发 ipv4_lpm 表项，实现基本转发功能
    """
    table_entry = p4info_helper.buildTableEntry(
        table_name="MyIngress.ipv4_lpm",
        match_fields={
            "hdr.ipv4.dstAddr": (dst_ip_addr, 32),
        },
        action_name="MyIngress.ipv4_forward",
        action_params={
            "dstAddr": dst_mac_addr,
            "port": port
        }
    )
    sw.WriteTableEntry(table_entry)
    print(f"Installed ipv4_lpm rule on {sw.name}: {dst_ip_addr} -> {dst_mac_addr} (port {port})")

# 新增，老师提示的函数：下发 UDP ACL 规则
def writeAcl_udpRules(p4info_helper, sw, match_fields, priority):
    table_entry = p4info_helper.buildTableEntry(
        table_name="MyIngress.acl_udp_ternary",
        match_fields={
            "hdr.udp.dstPort": match_fields
        },
        action_name="MyIngress.drop",
        priority=priority
    )
    sw.WriteTableEntry(table_entry)
    print(f"Installed ACL UDP rule on {sw.name}: match {match_fields}")

# 新增，自写的 IP ACL 规则（结构模仿上面的）
def writeAcl_ipRules(p4info_helper, sw, match_fields, priority):
    table_entry = p4info_helper.buildTableEntry(
        table_name="MyIngress.acl_ip_ternary",
        match_fields={
            "hdr.ipv4.dstAddr": match_fields
        },
        action_name="MyIngress.drop",
        priority=priority
    )
    sw.WriteTableEntry(table_entry)
    print(f"Installed ACL IP rule on {sw.name}: match {match_fields}")

def main():
    # 加载 p4info 文件和 bmv2 json 文件
    p4info_file_path = "build/acl.p4info.txt"
    bmv2_file_path = "build/acl.json"

    # 创建 p4info_helper
    p4info_helper = p4runtime_lib.helper.P4InfoHelper(p4info_file_path)

    try:
        # 创建与交换机的连接
        s1 = p4runtime_lib.bmv2.Bmv2SwitchConnection(
            name='s1',
            address='127.0.0.1:50051',
            device_id=0,
            proto_dump_file='logs/s1-p4runtime-requests.log'
        )

        # 设置 pipeline config
        s1.SetForwardingPipelineConfig(p4info=p4info_helper.p4info, bmv2_json_file_path=bmv2_file_path)
        print("Installed P4 Program on s1")

        # 安装 ipv4_lpm 表项，实现 h1, h2, h3 的正常转发
        write_ipv4_lpm_rule(p4info_helper, s1, "10.0.1.1", "00:00:00:00:01:01", 1)
        write_ipv4_lpm_rule(p4info_helper, s1, "10.0.1.2", "00:00:00:00:01:02", 2)
        write_ipv4_lpm_rule(p4info_helper, s1, "10.0.1.3", "00:00:00:00:01:03", 3)

        # 安装 ACL 表项：新的方式
        # 注意 match_fields 参数传的是 [值, 掩码]，priority 是优先级
        writeAcl_ipRules(p4info_helper, s1, ["10.0.1.4", 4294967295], 1)
        writeAcl_udpRules(p4info_helper, s1, [80, 65535], 1)

        print("All rules have been installed successfully.")

    except KeyboardInterrupt:
        print(" Shutting down.")
    except grpc.RpcError as e:
        print(f"gRPC Error: {e.details()} ({e.code().name})")

    finally:
        # 断开连接
        s1.shutdown()

if __name__ == '__main__':
    main()
